import tensorflow as tf
import function

# 获取对应卷积层
def vgg_layers(layer_names):
    # 加载已经在 imagenet 数据上预训练的 VGG 
    vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')
    vgg.trainable = False
    outputs = [vgg.get_layer(name).output for name in layer_names]
    model = tf.keras.Model([vgg.input], outputs)
    return model

# 定义模型类
class TransModel(tf.keras.models.Model):
    # 构造函数
    def __init__(self, style_layers, content_layers):
        super(TransModel, self).__init__()
        # 获取卷积核
        self.vgg =  vgg_layers(style_layers + content_layers)
        self.style_layers = style_layers
        self.content_layers = content_layers
        self.num_style_layers = len(style_layers)
        self.vgg.trainable = False

    def call(self, input):
        # 输入为归一化的数据，进行归一化解除，用于调用preprocess_input函数来处理图像的白化
        input = input * 255.0
        preprocessed_input = tf.keras.applications.vgg19.preprocess_input(input)
        output = self.vgg(preprocessed_input)
        style_output, content_output = (output[:self.num_style_layers], output[self.num_style_layers:])    
        # 对风格图需要计算gram矩阵，而内容图不需要
        style_output = [function.gram_matrix(style_output) for style_output in style_output]   
        # 转换为字典用于返回
        content_dic = {content_name:value for content_name, value in zip(self.content_layers, content_output)} 
        style_dic = {style_name:value for style_name, value in zip(self.style_layers, style_output)}

        return {'content':content_dic, 'style':style_dic}